import time
import cv2
from pyzed import sl
import numpy as np
import os
import argparse


def main(save_root, svo_file, rate):
    assert os.path.isfile(svo_file), svo_file
    os.makedirs(save_root, exist_ok=True)
    zed = sl.Camera()
    init_params = sl.InitParameters()
    init_params.set_from_svo_file(svo_file)
    init_params.svo_real_time_mode = False
    if zed.open(init_params) != sl.ERROR_CODE.SUCCESS:
        print("打开ZED相机失败"); exit(1)
    image_zed = sl.Mat()
    c = 0
    while True:
        c += 1
        if zed.grab() == sl.ERROR_CODE.SUCCESS:
            # if c<250:
            #     continue
            if np.random.rand()>rate:
                continue
            zed.retrieve_image(image_zed, sl.VIEW.LEFT)
            frame_id = str(zed.get_timestamp(sl.TIME_REFERENCE.IMAGE).get_nanoseconds())
            frame_bgr = image_zed.get_data()[::-1, ::-1, :3].copy()
            save_path = os.path.join(save_root,f'image_{frame_id}.png')
            cv2.imwrite(save_path, frame_bgr)
        else:
            break
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ZED Tennis Ball Tracker (Enhanced)")
    parser.add_argument("--save-root", type=str, help="Run in interactive parameter tuning mode.")
    parser.add_argument("--file", type=str, default=None, help="Path to an SVO file to use as input.")
    parser.add_argument("--rate", type=float, default=1.0)
    args = parser.parse_args()
    main(save_root=args.save_root, svo_file = args.file, rate=args.rate)